"""
3.	使用Python和OpenCV，借助背景减除算法进行运动车辆前景检测（14分）
"""
import numpy as np
import cv2 as cv
import os
import glob

# parameters
HIS = 500
ALPHA = 1e-3
FPS = 25
INTERVAL = 1000 // FPS

VER_INT = np.random.randint(0, np.iinfo(np.int32).max)
VER = f'v{VER_INT}'
print('VER', VER)
BASE_DIR, FILE_NAME = os.path.split(__file__)
path = 'data/traffic_1min.mp4'
VIDEO_PATH = os.path.join(BASE_DIR, path)
SAVE_DIR = os.path.join(BASE_DIR, '_save', FILE_NAME, VER)
print(f'We will save to {SAVE_DIR} ...')
os.makedirs(SAVE_DIR, exist_ok=True)
VIDEO_OUT_PATH = os.path.join(SAVE_DIR, 'result.avi')

# ①	读取视频文件traffic_1min.mp4，提取背景帧
print('Training to get background ...')
video = cv.VideoCapture(VIDEO_PATH)
bg = cv.createBackgroundSubtractorMOG2(HIS, detectShadows=False)
cnt = 0
while True:
    ret, img = video.read()
    if not ret:
        break
    cnt += 1
    if cnt > HIS:
        break
    bg.apply(img, None, ALPHA)
    if cnt % 25 == 0:
        print(f'{cnt}/{HIS} over.')
video.release()
print('Trained')
video_size = img.shape[:2][::-1]

# ②	使用形态学处理，比如背景消除、做帧差、膨胀腐蚀等，并将提取的每帧数据保存在文件夹里
video = cv.VideoCapture(VIDEO_PATH)
cnt = 0
kernel = np.ones((3, 3), dtype=np.uint8)
while True:
    ret, img = video.read()
    if not ret:
        print('Video over.')
        break
    cnt += 1
    if 27 == 0xFF & cv.waitKey(INTERVAL):
        print('Interrupted by user.')
        break
    mask = bg.apply(img, None, 0.)
    mask = cv.morphologyEx(mask, cv.MORPH_OPEN, kernel)
    mask = cv.morphologyEx(mask, cv.MORPH_CLOSE, kernel)
    mask = cv.dilate(mask, kernel, iterations=2)
    img = cv.bitwise_and(img, img, mask=mask)

    cv.imshow('Traffic', img)
    file_path = os.path.join(SAVE_DIR, f'img_{cnt:010d}.jpg')
    cv.imwrite(file_path, img)
video.release()

# ③	文件夹里的所有图像帧，保存为视频文件并可播放
print(f'Writing video into {VIDEO_OUT_PATH} ...')
fourcc = cv.VideoWriter_fourcc(*list('XVID'))
video = cv.VideoWriter(VIDEO_OUT_PATH, fourcc, FPS, video_size)
file_list = glob.glob(SAVE_DIR + '/*.jpg')
file_names = [os.path.split(path)[1] for path in file_list]
file_names = sorted(file_names)
for file_name in file_names:
    file_path = os.path.join(SAVE_DIR, file_name)
    img = cv.imread(file_path, cv.IMREAD_COLOR)
    video.write(img)
video.release()
print('Video written.')

print('Over')
